import numpy as np


class BodyModel:
    def __init__(self, mapping, pairs) -> None:
        self.mapping = mapping
        self.pairs = pairs


BODY18 = BodyModel(
    mapping=[
        "nose",
        "left_eye",
        "right_eye",
        "left_ear",
        "right_ear",
        "left_shoulder",
        "right_shoulder",
        "left_elbow",
        "right_elbow",
        "left_wrist",
        "right_wrist",
        "left_hip",
        "right_hip",
        "left_knee",
        "right_knee",
        "left_ankle",
        "right_ankle",
        "neck",
    ],
    pairs=[
        [15, 13],
        [13, 11],
        [16, 14],
        [14, 12],
        [11, 12],
        [5, 7],
        [6, 8],
        [7, 9],
        [8, 10],
        [1, 2],
        [0, 1],
        [0, 2],
        [1, 3],
        [2, 4],
        [3, 5],
        [4, 6],
        [17, 0],
        [17, 5],
        [17, 6],
        [17, 11],
        [17, 12],
    ],
)

BODY18_FLAT = BodyModel(
    mapping=[
        "nose_x",
        "nose_y",
        "left_eye_x",
        "left_eye_y",
        "right_eye_x",
        "right_eye_y",
        "left_ear_x",
        "left_ear_y",
        "right_ear_x",
        "right_ear_y",
        "left_shoulder_x",
        "left_shoulder_y",
        "right_shoulder_x",
        "right_shoulder_y",
        "left_elbow_x",
        "left_elbow_y",
        "right_elbow_x",
        "right_elbow_y",
        "left_wrist_x",
        "left_wrist_y",
        "right_wrist_x",
        "right_wrist_y",
        "left_hip_x",
        "left_hip_y",
        "right_hip_x",
        "right_hip_y",
        "left_knee_x",
        "left_knee_y",
        "right_knee_x",
        "right_knee_y",
        "left_ankle_x",
        "left_ankle_y",
        "right_ankle_x",
        "right_ankle_y",
        "neck_x",
        "neck_y",
    ],
    pairs=[
        [(30, 31), (26, 27)],
        [(26, 27), (22, 23)],
        [(32, 33), (28, 29)],
        [(28, 29), (24, 25)],
        [(22, 23), (24, 25)],
        [(10, 11), (14, 15)],
        [(12, 13), (16, 17)],
        [(14, 15), (18, 19)],
        [(16, 17), (20, 21)],
        [(2, 3), (4, 5)],
        [(0, 1), (2, 3)],
        [(0, 1), (4, 5)],
        [(2, 3), (6, 7)],
        [(4, 5), (8, 9)],
        [(6, 7), (10, 11)],
        [(8, 9), (12, 13)],
        [(34, 35), (0, 1)],
        [(34, 35), (10, 11)],
        [(34, 35), (12, 13)],
        [(34, 35), (22, 23)],
        [(34, 35), (24, 25)],
    ],
)

BODY25 = BodyModel(
    mapping=[
        "nose",
        "neck",
        "right_shoulder",
        "right_elbow",
        "right_wrist",
        "left_shoulder",
        "left_elbow",
        "left_wrist",
        "mid_hip",
        "right_hip",
        "right_knee",
        "right_ankle",
        "left_hip",
        "left_knee",
        "left_ankle",
        "right_eye",
        "left_eye",
        "right_ear",
        "left_ear",
        "left_bigtoe",
        "left_smalltoe",
        "left_heel",
        "right_bigtoe",
        "right_smalltoe",
        "right_heel",
    ],
    pairs=[
        [1, 8],
        [1, 2],
        [1, 5],
        [2, 3],
        [3, 4],
        [5, 6],
        [6, 7],
        [8, 9],
        [9, 10],
        [10, 11],
        [8, 12],
        [12, 13],
        [13, 14],
        [1, 0],
        [0, 15],
        [15, 17],
        [0, 16],
        [16, 18],
        [2, 17],
        [5, 18],
        [14, 19],
        [19, 20],
        [14, 21],
        [11, 22],
        [22, 23],
        [11, 24],
    ],
)
""" #BODY25 annotated
    pairs_annotated={
        "Torso":[1, 8], 
        "Shoulder (right)":[1, 2],
        "Shoulder (left)":[1, 5],
        "Arm (right)":[2, 3],
        "Forearm (right)":[3, 4],
        "Arm (left)":[5, 6],
        "Forearm (left)":[6, 7],
        "Hip (right)":[8, 9],
        "Thigh (right)":[9, 10],
        "Leg (right)":[10, 11],
        "Hip (left)":[8, 12],
        "Thigh (left)":[12, 13],
        "Leg (left)":[13, 14],
        "Neck":[1, 0],
        "Eye (right)":[0, 15],
        "Ear (right)":[15, 17],
        "Eye (left)":[0, 16],
        "Ear (left)":[16, 18],
        "Foot (left)":[14, 19],
        "Toe (left)":[19, 20],
        "Heel (left)":[14, 21],
        "Foot (right)":[11, 22],
        "Toe (right)":[22, 23],
        "Heel (right)":[11, 24],
    }
"""

BODY25_FLAT = BodyModel(
    mapping=[
        "nose_x",
        "nose_y",
        "neck_x",
        "neck_y",
        "right_shoulder_x",
        "right_shoulder_y",
        "right_elbow_x",
        "right_elbow_y",
        "right_wrist_x",
        "right_wrist_y",
        "left_shoulder_x",
        "left_shoulder_y",
        "left_elbow_x",
        "left_elbow_y",
        "left_wrist_x",
        "left_wrist_y",
        "mid_hip_x",
        "mid_hip_y",
        "right_hip_x",
        "right_hip_y",
        "right_knee_x",
        "right_knee_y",
        "right_ankle_x",
        "right_ankle_y",
        "left_hip_x",
        "left_hip_y",
        "left_knee_x",
        "left_knee_y",
        "left_ankle_x",
        "left_ankle_y",
        "right_eye_x",
        "right_eye_y",
        "left_eye_x",
        "left_eye_y",
        "right_ear_x",
        "right_ear_y",
        "left_ear_x",
        "left_ear_y",
        "left_bigtoe_x",
        "left_bigtoe_y",
        "left_smalltoe_x",
        "left_smalltoe_y",
        "left_heel_x",
        "left_heel_y",
        "right_bigtoe_x",
        "right_bigtoe_y",
        "right_smalltoe_x",
        "right_smalltoe_y",
        "right_heel_x",
        "right_heel_y",
    ],
    pairs=[
        [(2, 3), (16, 17)],
        [(2, 3), (4, 5)],
        [(2, 3), (10, 11)],
        [(4, 5), (6, 7)],
        [(6, 7), (8, 9)],
        [(10, 11), (12, 13)],
        [(12, 13), (14, 15)],
        [(16, 17), (18, 19)],
        [(18, 19), (20, 21)],
        [(20, 21), (22, 23)],
        [(16, 17), (24, 25)],
        [(24, 25), (26, 27)],
        [(26, 27), (28, 29)],
        [(2, 3), (0, 1)],
        [(0, 1), (30, 31)],
        [(30, 31), (34, 35)],
        [(0, 1), (32, 33)],
        [(32, 33), (36, 37)],
        [(4, 5), (34, 35)],
        [(10, 11), (36, 37)],
        [(28, 29), (38, 39)],
        [(38, 39), (40, 41)],
        [(28, 29), (42, 43)],
        [(22, 23), (44, 45)],
        [(44, 45), (46, 47)],
        [(22, 23), (48, 49)],
    ],
)

# fmt: off
BODY25_to_BODY18_indices = [0, 16, 15, 18, 17, 5, 2, 6, 3, 7, 4, 12, 9, 13, 10, 14, 11, 1]
BODY25flat_to_BODY18flat_indices = [0, 1, 32, 33, 30, 31, 36, 37, 34, 35, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15, 8, 9, 24, 25, 18, 19, 26, 27, 20, 21, 28, 29, 22, 23, 2, 3]
# fmt: on


def BODY25_to_BODY18(body25_keypoints: np.ndarray):
    assert body25_keypoints.shape == 25
    return body25_keypoints[BODY25_to_BODY18_indices]
